import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 网格分辨率参数
parm = [8, 16, 32, 64]

# 创建一个大图，包含多个子图
fig = plt.figure(figsize=(15, 10))

# 遍历每个网格分辨率
for idx, resolution in enumerate(parm):
    # 读取数据
    filename = f'Problem2_parm{resolution}.txt'
    data = np.loadtxt(filename)

    # 提取x, y, z坐标
    x = data[:, 0]
    y = data[:, 1]
    z = data[:, 2]

    # 创建网格
    m = int(np.sqrt(len(x)))  # 网格分辨率
    x = x.reshape((m, m))
    y = y.reshape((m, m))
    z = z.reshape((m, m))

    # 添加子图
    ax = fig.add_subplot(2, 2, idx + 1, projection='3d')  # 2行2列的子图布局

    # 绘制表面图
    surf = ax.plot_surface(x, y, z, cmap='viridis')

    # 添加颜色条
    fig.colorbar(surf, ax=ax, shrink=0.6, aspect=10)

    # 设置标签和标题
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('U')
    ax.set_title(f'Grid Resolution: {resolution}x{resolution}')

# 调整布局
plt.tight_layout()
plt.savefig('../../figure/F1_Square_M.png')
# 显示图形
plt.show()